import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
import json
import time
import os
from PIL import Image
import torch.nn as nn

import pretrainedmodels
import torchattacks
from torchattacks import *
import argparse

from all_in_one_imagenet_ILA_ATA import get_source_layers
from attacks_vgg import ifgsm, EILA
from att_utils import imshow, image_folder_custom_label

# Training settings
parser = argparse.ArgumentParser(description='Demo for ILA on ImageNet')
parser.add_argument('--modeltype', type=str, default='resnet18',
                    help='ResNet18 | DenseNet121 | alexnet | SqueezeNet1.0')
parser.add_argument('--layerindex1', type=int, default='9', help='layer index to emphasize with ila projection')
parser.add_argument('--layerindex2', type=int, default='9', help='layer index to emphasize with ila projection')
parser.add_argument('--layerindex3', type=int, default='9', help='layer index to emphasize with ila projection')
# parser.add_argument('--imagepath', type=str, default='eft_image_test.jpg', help='path to image to test')
# parser.add_argument('--outpath', type=str, default='adv_example.jpg', help='path for output image')
# parser.add_argument('--imagelabel', type=int, default=27, help='imagenet label (0-999)')
# parser.add_argument('--niters_baseline', type=int, default=20, help='number of iterations of baseline attack')
parser.add_argument('--niters_ila', type=int, default=10, help='number of iterations of ILA')
parser.add_argument('--epsilon', type=float, default=16 / 255,
                    help='epsilon on 0..1 range, 0.03 corresponds to ~8 in the imagenet scale')
parser.add_argument('--alpha', type=float, default=2 / 255,
                    help='alpha on 0..1 range, 0.03 corresponds to ~8 in the imagenet scale')
opt = parser.parse_args()


def save_image_tensor(input_tensor: torch.Tensor, filename):
    """
    将tensor保存为图片
    :param input_tensor: 要保存的tensor
    :param filename: 保存的文件名
    """
    assert (len(input_tensor.shape) == 4 and input_tensor.shape[0] == 1)
    # 复制一份
    input_tensor = input_tensor.clone().detach()
    # 到cpu
    input_tensor = input_tensor.to(torch.device('cpu'))
    # 反归一化
    # input_tensor = unnormalize(input_tensor)
    # input_tensor = ((input_tensor * std) + mean)
    # unorm = UnNormalize(mean=mean, std=std)
    # unorm(input_tensor)
    torchvision.utils.save_image(input_tensor, filename)


class Normalize(nn.Module):
    def __init__(self, mean, std):
        super(Normalize, self).__init__()
        self.register_buffer('mean', torch.Tensor(mean))
        self.register_buffer('std', torch.Tensor(std))

    def forward(self, input):
        # Broadcasting
        mean = self.mean.reshape(1, 3, 1, 1)
        std = self.std.reshape(1, 3, 1, 1)
        return (input - mean) / std


if __name__ == '__main__':
    class_idx = json.load(
        open("/mnt/zhanrui/Yangshanshan/adversarial-attacks-pytorch/sample/data/imagenet_class_index.json"))
    idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))]

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),  # ToTensor : [0, 255] -> [0, 1]
        # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    imagnet_data = image_folder_custom_label(
        root='/mnt/zhanrui/Yangshanshan/adversarial-attacks-pytorch/sample/data/imagenet/val/', transform=transform,
        idx2label=idx2label)
    data_loader = torch.utils.data.DataLoader(imagnet_data, batch_size=1, shuffle=False)

    norm_layer = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    model = nn.Sequential(
        norm_layer,
        pretrainedmodels.__dict__['vgg19'](num_classes=1000, pretrained='imagenet')
    )

    # model = pretrainedmodels.__dict__['resnet50'](num_classes=1000, pretrained='imagenet')

    model = model.eval().cuda()
    # for name, _ in model.named_modules():
    #     print(name)

    # 准备工作
    attack = BIM(model, eps=16 / 255, alpha=2 / 255, steps=10)
    # attack = MIFGSM(model, eps=16 / 255, alpha=2 / 255, steps=10, decay=0.1)
    # attack = DIFGSM(model, eps=16 / 255, alpha=2 / 255, steps=10, diversity_prob=0.5, resize_rate=0.9)
    save_path = "/mnt/zhanrui/Yangshanshan/adversarial-attacks-pytorch/sample/data/imagenet/3-adv_bimIAAl8_vgg19_ens3"
    correct = 0
    step = 0
    total = 0
    start = time.time()
    # 加载图像并进行攻击
    for images, labels in data_loader:
        images = images.cuda()
        source_layers = get_source_layers(opt.modeltype, model)
        # ifgsm_guide = ifgsm(model, images, labels, learning_rate=0.008, epsilon=opt.epsilon, niters=opt.niters_ila,
        #                     dataset='imagenet')
        ifgsm_guide = attack(images, labels)
        img_ila = EILA(model, images, ifgsm_guide, labels, source_layers[opt.layerindex1][1][1],
                       source_layers[opt.layerindex2][1][1], source_layers[opt.layerindex3][1][1], alpha=opt.alpha,
                       epsilon=opt.epsilon, niters=opt.niters_ila, dataset='imagenet')
        labels = labels.cuda()
        outputs = model(img_ila)

        _, pre = torch.max(outputs.data, 1)

        step += 1
        total += labels.size(0)
        name = imagnet_data.samples[step - 1][0]
        # img_pil = Image.open(name)
        # img_pil_resize = transform_resize(img_pil.copy())
        # 注意起始位置：val是19，val_test是24
        file_name = name[78:-28]
        both_name = name[-28:]
        new_path = save_path + file_name
        new_name = new_path + both_name
        if not os.path.exists(new_path):
            os.makedirs(new_path)

        # correct += (pre == labels).sum()
        correct += (pre == labels.cuda()).sum()

        # save image my method
        # save_image_tensor(img_ila, new_name, mean=mean, std=stddev)
        save_image_tensor(img_ila, new_name)
        # imshow(img_ila.cpu().data, [imagnet_data.classes[i] for i in pre])

    print('Total elapsed time (sec): %.2f' % (time.time() - start))
    print('Robust accuracy: %.2f %%' % (100 * float(correct) / total))
